import cv2
import matplotlib.pyplot as plt
from ultralytics import YOLO
import os


class VehicleDetector:
    def __init__(self, model_path='yolov8n.pt'):
        """
        车辆检测系统初始化

        参数:
        model_path: YOLO模型路径（默认使用YOLOv8n预训练模型）
        """
        # 加载预训练的YOLOv8模型
        self.model = YOLO(model_path)

        # 定义车辆类别ID和对应的英文名称
        self.vehicle_classes = {
            2: "car",
            5: "bus",
            7: "truck"
        }

    def detect_vehicles(self, image_path=None, image=None):
        # 图像加载
        if image is None:
            if image_path is None:
                raise ValueError("必须提供image_path或image参数")
            img = cv2.imread(image_path)
            if img is None:
                raise FileNotFoundError(f"无法读取图像: {image_path}")
        else:
            img = image

        # 转换为RGB格式
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # 执行目标检测
        results = self.model(img_rgb)

        # 处理检测结果
        vehicle_count = 0
        result_img = img.copy()
        corner_points = []  # 存储四个角点坐标

        for box in results[0].boxes:
            class_id = int(box.cls)
            if class_id in self.vehicle_classes:
                # 车辆计数增加
                vehicle_count += 1

                # 获取边界框坐标
                x1, y1, x2, y2 = map(int, box.xyxy[0])

                # 计算四个角点坐标
                top_left = (x1, y1)
                top_right = (x2, y1)
                bottom_left = (x1, y2)
                bottom_right = (x2, y2)
                corner_points.append({
                    'top_left': top_left,
                    'top_right': top_right,
                    'bottom_left': bottom_left,
                    'bottom_right': bottom_right,
                    'class_id': class_id,
                    'class_name': self.vehicle_classes[class_id],
                    'confidence': float(box.conf)
                })

                # 获取置信度分数
                confidence = float(box.conf)

                # 获取类别名称
                class_name = self.vehicle_classes[class_id]

                # 绘制边界框和标签
                cv2.rectangle(result_img, (x1, y1), (x2, y2), (0, 255, 0), 2)  # 绿色边界框
                label = f"{class_name}: {confidence:.2f}"
                cv2.putText(result_img, label, (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)  # 绿色文本

                # 在四个角点绘制小圆圈（红色）
                cv2.circle(result_img, top_left, 5, (0, 0, 255), -1)
                cv2.circle(result_img, top_right, 5, (0, 0, 255), -1)
                cv2.circle(result_img, bottom_left, 5, (0, 0, 255), -1)
                cv2.circle(result_img, bottom_right, 5, (0, 0, 255), -1)

        return result_img, vehicle_count, corner_points

    def detect_from_camera(self, camera_id=0):
        """
        从摄像头实时检测车辆

        参数:
        camera_id: 摄像头设备ID
        """
        # 初始化摄像头
        cap = cv2.VideoCapture(camera_id)

        if not cap.isOpened():
            raise ValueError(f"无法打开摄像头 {camera_id}")

        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    print("获取帧失败")
                    break

                # 检测车辆
                result_img, vehicle_count, corner_points = self.detect_vehicles(image=frame)

                # 显示车辆计数
                cv2.putText(result_img, f"quentity: {vehicle_count}", (10, 30),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)  # 红色文本

                # 打印角点坐标到控制台
                for i, points in enumerate(corner_points):
                    print(f"车辆 {i+1}: {points['class_name']} ({points['confidence']:.2f})")
                    print(f"  左上角: {points['top_left']}")
                    print(f"  右上角: {points['top_right']}")
                    print(f"  左下角: {points['bottom_left']}")
                    print(f"  右下角: {points['bottom_right']}")

                # 显示结果
                cv2.imshow('车辆检测系统', result_img)

                # 按ESC键退出
                key = cv2.waitKey(1)
                if key == 27:  # ESC键
                    break
        finally:
            cap.release()
            cv2.destroyAllWindows()


def main():
    # 初始化车辆检测器
    detector = VehicleDetector()

    #从图像文件检测
    image_path = r'' # 图像路径

    if os.path.exists(image_path):
        try:
            # 检测车辆
            result_img, vehicle_count, corner_points = detector.detect_vehicles(image_path)

            # 显示结果
            plt.figure(figsize=(10, 8))
            plt.imshow(cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB))
            plt.title(f"quentity: {vehicle_count}")
            plt.axis('off')
            plt.show()

            print(f"图像中检测到 {vehicle_count} 辆车辆")
            for i, points in enumerate(corner_points):
                print(f"车辆 {i+1}: {points['class_name']} ({points['confidence']:.2f})")
                print(f"  左上角: {points['top_left']}")
                print(f"  右上角: {points['top_right']}")
                print(f"  左下角: {points['bottom_left']}")
                print(f"  右下角: {points['bottom_right']}")
        except Exception as e:
            print(f"图像处理错误: {e}")
    else:
        print(f"未找到图像文件: {image_path}")
        print("正在切换到摄像头模式...")

        # 从摄像头检测
        try:
            detector.detect_from_camera()
        except Exception as e:
            print(f"摄像头启动错误: {e}")
7
if __name__ == "__main__":
    main()